# -*- coding: utf-8 -*-
"""
PBAR_ExamineZspecGainDatasetsSet2.py

2/14/2014, JK
"""


import PBAR_Zspec
import numpy as np
import matplotlib.pyplot as plt
import copy, csv, os
from scipy.optimize import curve_fit
from scipy.stats import tmean
# PMT curve fit
def PMT_Curve(x, a, b):
    return(a* np.power(x, b))

# List of good and bad detectors
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])


badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,62,76,80,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136])


goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecList -= 1
goodZspecList -= 1

baseDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3'


# create list of dataset names
# not including ea91 because bad
datasetNameList = np.array(['ec%d'%i for i in xrange(49, 61)])
filenameList = np.array([(i+'.csv') for i in datasetNameList])
fullFilenameList = np.array([os.path.join(baseDir,i+'.csv') for i in datasetNameList])
dat = PBAR_Zspec.ReadZspec(fullFilenameList)

# Important dataset information
#voltageOffset = np.array([0, -9, -6, 0, 9, 21, 36, -3, 0, 3, 6, 12, 0, 5])
voltageOffset = np.array([0, 5, 10, 15, 20, 0, -5, -15, -30, -20, -10, 0])

#epochTimeList = np.array([1378493600, 1378494983, 1378495716, 1378496415, 1378497181, 1378498125, 1378498872, 1378500092, 1378500876,\
#    1378501619, 1378503186, 1378503992, 1378505158, 1378507618])
epochTimeList = np.array([1392337140, 1392337320, 1392337680, 1392337860, 1392338400, 1392338700, \
    1392339120, 1392339540, 1392336360, 1392337800, 1392337980, 1392338460 ])

epochTimeList0 = epochTimeList - epochTimeList[0]

# read in voltage
voltageDat = np.genfromtxt(r'C:\Users\jkwong\Documents\Work\PBAR\data3\voltage_20140213_161305.txt', delimiter = ',')
voltageSet = voltageDat[:,4]
voltageMeasured = voltageDat[:,5]


acqTime = 120.0
pulseRate = 60.
binBaseline = 100.
rateRange = [0.5, 2.0] 
rateLook = 1.0
binBounds = [20, 200]

gainShiftAll = np.zeros((len(dat), dat[0].shape[1]))
gainShift0All = np.zeros((len(dat), dat[0].shape[1]))

for i in xrange(len(dat)):
    #    print dat[0].shape
    binArray = np.arange(dat[i].shape[0])
    
    # fit an exponential
    #    pfitAll = np.zeros((dat[0].shape[1], 2))
    binMeanAll = np.zeros(dat[i].shape[1])
    
    for ch in xrange(dat[i].shape[1]):
        cutt = (binArray > binBounds[0]) & (binArray < binBounds[1]) & \
        (dat[i][:,ch]/acqTime > rateRange[0]) & (dat[i][:,ch]/acqTime < rateRange[1])
        if any(cutt):
            temp = np.polyfit(binArray[cutt], np.log(dat[i][cutt,ch]/acqTime),1)
            pfit0 = [temp[0], np.exp(temp[1])]
            binMean = (1.0/pfit0[0]) * np.log(rateLook/pfit0[1])
#            print 'a'
        else:
            binMean = 0.0
        binMeanAll[ch] = binMean
    #        print i+1, binMean
    gainShift0All[i,:] = binMeanAll / binBaseline

for i in xrange(gainShift0All.shape[0]):
    gainShiftAll[i,:]= gainShift0All[i,:] / gainShift0All[0,:]

#
## # we don't need this as the temperature was quite constant
## First we have to remove the gain shift due to temperature buy using the deltaV = 0 datasets
#
#gainShiftCorrectedAll =  copy.copy(gainShiftAll)
## apply correction gain shift due to temperature
#cut = voltageOffset == 0    
#x = epochTimeList0[cut]
#y = gainShiftAll[cut,:]
#
## Correct each channel separately
#pfitGainShiftVsTimeList = np.zeros((136, 2))
#for i in xrange(136):
#    pfit = np.polyfit(x, np.log(y[:,i]), 1)
#    pfitGainShiftVsTimeList[i,:] = pfit
#    # correct gain
#    gainShiftCorrectedAll[:,i] = gainShiftAll[:,i] * np.exp(np.polyval(pfit, 0)) / np.exp(np.polyval(pfit, epochTimeList0))


gainShiftCorrectedAll = gainShiftAll

# Fit PMT voltage gain curves, one PMT at a time
x = voltageOffset
y = gainShiftCorrectedAll

pfitGainShiftVsVoltageList = np.zeros((136, 2))
poptGainShiftVsVoltageList = np.zeros((136, 2))
for i in xrange(136):
    cut = y[:,i] > 0.0
    try:
        pfit = np.polyfit(x[cut], log(y[cut,i]), 1)
    except:
        pfit = np.array([-1, -1])
    pfitGainShiftVsVoltageList[i,:] = pfit
    try:
        startingParam = [2., 20.]
        popt, pcov = curve_fit(PMT_Curve, (x[cut] + voltageSet[i])/1000.0, y[cut,i], p0 = startingParam)
#        poptGainShiftVsVoltageList[i,0] = popt[0]
        poptGainShiftVsVoltageList[i,0] = popt[0] / 1000.0 ** popt[1]     
        poptGainShiftVsVoltageList[i,1] = popt[1]
    except:
        poptGainShiftVsVoltageList[i,:] = np.array([-1,-1])

gainCorrectionCoefficient = 1/pfitGainShiftVsVoltageList[:,0]
# correct the channels that have nan values
cutt = ~isnan(gainCorrectionCoefficient) & ~isinf(gainCorrectionCoefficient) & (gainCorrectionCoefficient > 0)
gainCorrectionCoefficientAverage = gainCorrectionCoefficient[cutt].mean()
gainCorrectionCoefficient[~cutt] = gainCorrectionCoefficientAverage

with open(os.path.join(baseDir, 'GainCalibrationCoefficientsSet2.txt'), 'wb') as fid:
    csvWriterObj = csv.writer(fid)
    csvWriterObj.writerow(gainCorrectionCoefficient)
    fid.close()

# correct the other way -  Idon't think that these values are used 2/14/2014
gainCorrectionCoefficient2 = poptGainShiftVsVoltageList[:,1]
cutt = ~isnan(gainCorrectionCoefficient2) & ~isinf(gainCorrectionCoefficient2) & (gainCorrectionCoefficient2 > 0)
gainCorrectionCoefficient2Average = gainCorrectionCoefficient2[cutt].mean()
gainCorrectionCoefficient2[~cutt] = gainCorrectionCoefficient2Average

with open(os.path.join(baseDir, 'GainCalibrationCoefficients2Set2.txt'), 'wb') as fid:
    csvWriterObj = csv.writer(fid)
    csvWriterObj.writerow(gainCorrectionCoefficient2)
    fid.close()


# PLOTS

# VERSUS TIME

# plot gainshift versus time for delta v = 0 
cut = voltageOffset == 0
plt.figure()
plt.grid()
detectorIndex = 57
plot(epochTimeList[cut] - epochTimeList[cut][0], gainShiftAll[cut,detectorIndex], 'xk', markersize = 12, mew = 2)

plt.xlabel('Time Since First Dataset (sec)')
plt.ylabel('Gain Shift')
plt.title('Detector #%d' %(detectorIndex+1))

# plot gainshift versus time for delta v = 0 
plt.figure()
detectorIndex = 60
plt.plot(epochTimeList, gainShiftAll[:,detectorIndex], '.k')

plt.xlabel('Voltage Offset')
plt.ylabel('Gain Shift')
plt.title('Detector #%d' %i)

# GAIN SHIFT VERSUS TIME FOR V = 0
plt.figure()
plt.grid()
cut = voltageOffset == 0
y = gainShift0All[cut,:]
epochTimeListShortList = epochTimeList0[cut]

for i in xrange(y.shape[0]):
    cut2 = ~isnan(y[i,:])
    plt.plot(np.arange(1,137)[cut2], y[i,cut2]/y[0,cut2], marker = 'x', label = '%d sec' %epochTimeListShortList[i])

plt.xlabel('Detector Number')
plt.ylabel('Gain Shift')
plt.legend()
plt.ylim((0.7, 1.3))


# Gain correction parameters as function of detector number
plt.figure()
plt.grid()
cut = ~isnan(gainCorrectionCoefficient)
plot(np.arange(1, 137)[cut], gainCorrectionCoefficient[cut], 'xk')

cut = np.zeros(136) == 1
cut[goodZspecList] = True
cut = cut & ~isnan(gainCorrectionCoefficient)
plot(np.arange(1, 137)[cut], gainCorrectionCoefficient[cut], 'xr')

plt.xlabel('Detector Number', fontsize = 16)
plt.ylabel('Gain Correction Parameter', fontsize = 16)




# Gain correction parameters as function of detector number - GOOD detectors
plt.figure()
plt.grid()
#cut = ~isnan(gainCorrectionCoefficient)
#plot(np.arange(1, 137)[cut], gainCorrectionCoefficient[cut], 'xk')

cut = np.zeros(136) == 1
cut[goodZspecList] = True
cut = cut & ~isnan(gainCorrectionCoefficient)
plot(np.arange(1, 137)[cut], gainCorrectionCoefficient[cut], 'xk')

plt.xlabel('Detector Number', fontsize = 16)
plt.ylabel('Gain Correction Parameter', fontsize = 16)

plt.xlim((0, 136))



# Gain correction parameters as function of detector number
plt.figure()
plt.grid()
#cut = ~isnan(gainCorrectionCoefficient)
#plot(np.arange(1, 137)[cut], gainCorrectionCoefficient[cut], 'xk')

cut = np.zeros(136) == 1
cut[goodZspecList] = True
cut = cut & ~isnan(gainCorrectionCoefficient)
plot(np.arange(1, 137)[cut], gainCorrectionCoefficient2[cut], 'xk')

plt.xlabel('Detector Number', fontsize = 16)
plt.ylabel('Gain Correction Parameter', fontsize = 16)

plt.xlim((0, 136))








# Gainshift as a function of voltage; good detectors
plt.figure()
for i in xrange(136):
    if i in goodZspecList:
        plt.plot(voltageOffset, gainShiftAll[:,i])


# GAIN SHIFT VERSUS VOLTAGE WITH FITS

plt.figure()
plt.grid()

i = 70
#i =67

plt.plot(voltageOffset, gainShiftAll[:,i], '.k', label = 'Uncorrected', markersize = 15, alpha = 0.5)
plt.plot(voltageOffset, gainShiftCorrectedAll[:,i], 'xr', label = 'Corrected', markersize = 15, alpha = 0.5)

xarray = np.arange(-30,30)
yarray = np.exp(np.polyval(pfitGainShiftVsVoltageList[i,:], xarray))
plt.plot(xarray, yarray, label = \
    'Best fit to corrected, a*exp(b*v), a = %3.3e, b = %3.3f' %(exp(pfitGainShiftVsVoltageList[i,1]), pfitGainShiftVsVoltageList[i,0]))

plt.plot(xarray, PMT_Curve(xarray + voltageSet[i], poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]), \
    label = 'Best fit to corrected, a*v^b, a = %3.3e, b = %3.3f' %(poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]))

plt.xlabel('Voltage Offset', fontsize = 16)
plt.ylabel('Gain Shift', fontsize = 16)
plt.title('Detector #%d, V = %d' %(i+1, voltageSet[i]))
plt.legend()
plt.yscale('log')

# plot gainshift as a function of voltage, single detector

for i in goodZspecList:
    try:
        plt.figure()
        plt.grid()
    
        #i = 60
        plt.plot(voltageOffset, gainShiftAll[:,i], '.k', label = 'Uncorrected', markersize = 15, alpha = 0.5)
        plt.plot(voltageOffset, gainShiftCorrectedAll[:,i], 'xr', label = 'Corrected', markersize = 15, alpha = 0.5)
        
        xarray = np.arange(-30,30)
        yarray = np.exp(np.polyval(pfitGainShiftVsVoltageList[i,:], xarray))
        plt.plot(xarray, yarray, label = \
            'Best fit to corrected, a*exp(b*v), a = %3.3e, b = %3.3f' %(np.exp(pfitGainShiftVsVoltageList[i,1]), pfitGainShiftVsVoltageList[i,0]))
        
        plt.plot(xarray, PMT_Curve(xarray + voltageSet[i], poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]), \
            label = 'Best fit to corrected, a*v^b, a = %3.3e, b = %3.3f' %(poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]))
        
        plt.xlabel('Voltage Offset')
        plt.ylabel('Gain Shift')
        plt.title('Detector #%d, V = %d' %(i, voltageSet[i]))
        plt.legend()
        plt.yscale('log')
        savefig(os.path.join(baseDir, 'GainCalibration_%d') %(i+1))
        plt.close()
    except:
        print("Skipping")


# EXAMINE FIT PARAMETERS

# Examine the fit parameters
# fit 1 parameters

plt.figure()
plt.grid()
cut = goodZspecList
plot(np.arange(1,137), poptGainShiftVsVoltageList[:,1], '.k')
plot(np.arange(1,137)[cut], poptGainShiftVsVoltageList[cut,1], '.r')

# calculate the trimmed mean
cutt = poptGainShiftVsVoltageList[cut,1] > 0
plot([0, 136], poptGainShiftVsVoltageList[cut,1][cutt].mean() * np.ones(2), '-r')

plt.xlabel('Detector Number')
plt.ylabel("PMT Coefficient")

# Examine the fit parameters
# fit 2 parameters

figure()
plt.grid()

cut = goodZspecList
plot(np.arange(1,137), pfitGainShiftVsVoltageList[:,0], '.k')
plot(np.arange(1,137)[cut], pfitGainShiftVsVoltageList[cut,0], '.r')
plt.xlabel('Detector Number')
plt.ylabel("PMT Coefficient")
plt.ylim([0,0.04])


# Examine expected voltage increment for given gain shift
gs = 1.05
vIncrement1 = gainCorrectionCoefficient * log(gs)
vIncrement2 = (np.power(gs, 1/gainCorrectionCoefficient2) - 1.0) * voltageSet

plt.figure()
plot(np.arange(1, 137), vIncrement2, '.k')
cutt = goodZspecList
plot(np.arange(1, 137)[cutt], vIncrement2[cutt], '.r')

# SCATTER Plot of the two v increment estimates
plt.figure()
plt.grid()
plot(vIncrement1[goodZspecList], vIncrement2[goodZspecList], '.k')

plt.xlim([0, 3.5])
plt.ylim([0, 3.5])

plt.xlabel('V Increment 1')
plt.ylabel('V Increment 2')

# SCATTER Plot of the ratio of the increment estimates
plt.figure()
plt.grid()
plot(vIncrement1[goodZspecList], vIncrement2[goodZspecList]/vIncrement1[goodZspecList], '.k')

plt.xlim([0, 3.5])
plt.ylim([0, 3.5])

plt.xlabel('V Increment 1')
plt.ylabel('V Increment 2 / V Increment 1')

# Different in v-increment vs detector number
plt.figure()
plt.grid()
plot(np.arange(1, 137)[goodZspecList], vIncrement2[goodZspecList] - vIncrement1[goodZspecList], '.k')

#plt.xlim([0, 3.5])
#plt.ylim([0, 3.5])

plt.xlabel('Detector Number')
plt.ylabel('V increment 2 - V increment 1')